from matplotlib import pyplot as plt
import numpy as np
import os
from pyspark import SparkContext
from pyspark.mllib.classification import SVMWithSGD, NaiveBayes
from pyspark.mllib.linalg import DenseVector
from pyspark.mllib.regression import LabeledPoint

os.environ['PYSPARK_PYTHON'] = "python3"

blue_points = [(1.27, 2.2), (1.17, 1.63), (1.4, 2.8), (2.13, 3.0), (2.57, 2.87), (2.5, 2.1), (2.13, 1.73), (2.97, 2.43), (2.9, 2.67), (2.13, 2.7), (1.93, 2.37), (1.83, 2.43), (1.63, 1.97), (2.23, 2.0), (2.27, 2.4), (2.67, 2.53), (1.8, 2.67), (1.53, 1.47), (1.77, 1.67), (2.23, 1.57), (2.67, 1.73), (3.5, 2.47), (1.67, 1.2), (1.47, 0.93), (0.8, 0.93), (0.8, 1.7), (1.17, 2.6), (1.77, 2.13), (1.57, 1.53)]
red_points = [(1.17, 0.47), (1.7, 0.77), (2.9, 0.67), (3.9, 0.23), (4.9, 0.87), (4.67, 1.6), (4.57, 2.4), (4.33, 2.73), (3.53, 3.1), (3.53, 2.47), (3.53, 2.43), (3.83, 1.53), (3.1, 1.6), (3.3, 1.87), (2.8, 1.3), (2.43, 1.3), (2.43, 0.83), (1.8, 0.27), (2.2, 0.43), (3.1, 0.43), (4.37, 1.2), (3.53, 1.2), (3.4, 1.43), (3.4, 0.47), (4.07, 0.77), (4.07, 1.87), (4.07, 2.47), (3.3, 2.47), (2.7, 2.33), (2.57, 2.0), (2.23, 1.87), (2.23, 1.67), (2.9, 1.8), (3.03, 2.2), (3.17, 2.83), (4.0, 2.83), (4.1, 3.27), (4.7, 2.43), (3.63, 2.3), (2.97, 1.53), (1.73, 1.0), (2.17, 1.2), (2.17, 0.7), (3.77, 0.73)]

sc = SparkContext.getOrCreate()

label_map = {
    0: "RED",
    1: "BLUE"
}

labeled_points = [LabeledPoint(0, DenseVector((x, y))) for x, y in red_points] + [LabeledPoint(1, DenseVector((x, y))) for x, y in blue_points]

model = NaiveBayes.train(sc.parallelize(labeled_points))

pt = input("Enter point(x,y):").split(",")
X, Y = float(pt[0]), float(pt[1])

d = model.predict((X ,Y))
print(label_map[d])


plt.scatter([x for x, _ in red_points], [y for _, y in red_points], color="r")
plt.scatter([x for x, _ in blue_points], [y for _, y in blue_points], color="b")

x = np.arange(0, 5, 0.1)

# print(model.weights, model.intercept)
plt.ylim(0, 5)

plt.scatter([X], [Y], color='g')
# plt.plot(x, (model.weights[0])* x + model.weights[1] , color="g")
# plt.plot(x, model.weights[1] * x, color="g")

plt.show()

